#!/usr/bin/env python3

RATE = 8
A = 12
B = 6

def ascon_aead_encrypt(key, nonce, associated_data, plaintext):
    assert len(key) == 16 and len(nonce) == 16
    S = [0] * 5
    ascon_initialize(S, key, nonce)
    ascon_process_associated_data(S, associated_data)
    ciphertext = ascon_process_plaintext(S, plaintext)
    tag = ascon_finalize(S, key)
    return ciphertext + tag

def ascon_aead_decrypt(key, nonce, associated_data, ciphertext):
    assert len(key) == 16 and len(nonce) == 16 and len(ciphertext) >= 16
    S = [0] * 5
    ascon_initialize(S, key, nonce)
    ascon_process_associated_data(S, associated_data)
    plaintext = ascon_process_ciphertext(S, ciphertext[:-16])
    tag = ascon_finalize(S, key)
    if tag == ciphertext[-16:]:
        return plaintext
    else:
        return None

def ascon_initialize(S, key, nonce):
    k = len(key) * 8
    iv_zero_key_nonce = bytes([k, RATE * 8, A, B]) + (20-len(key)) * b"\x00" + key + nonce
    S[0], S[1], S[2], S[3], S[4] = bytes_to_state(iv_zero_key_nonce)
    ascon_permutation(S, A)
    zero_key = bytes_to_state((40-len(key)) * b"\x00" + key)
    for i in range(5):
        S[i] ^= zero_key[i]

def ascon_process_associated_data(S, associated_data):
    if associated_data:
        a_padding = b"\x80" + (RATE - (len(associated_data) % RATE) - 1) * b"\x00"
        a_padded = associated_data + a_padding
        for i in range(0, len(a_padded), RATE):
            S[0] ^= bytes_to_int(a_padded[i:i+8])
            if RATE == 16:
                S[1] ^= bytes_to_int(a_padded[i+8:i+16])
            ascon_permutation(S, B)
    S[4] ^= 1

def ascon_process_plaintext(S, plaintext):
    p_lastlen = len(plaintext) % RATE
    p_padding = b"\x80" + (RATE - p_lastlen - 1) * b"\x00"
    p_padded = plaintext + p_padding
    ciphertext = bytes()
    for i in range(0, len(p_padded) - RATE, RATE):
        S[0] ^= bytes_to_int(p_padded[i:i+8])
        ciphertext += int_to_bytes(S[0], 8)
        ascon_permutation(S, B)
    return ciphertext + int_to_bytes(S[0], 8)[:p_lastlen]

def ascon_process_ciphertext(S, ciphertext):
    c_lastlen = len(ciphertext) % RATE
    c_padded = ciphertext + (RATE - c_lastlen) * b"\x00"
    plaintext = bytes()
    for i in range(0, len(c_padded) - RATE, RATE):
        Ci = bytes_to_int(c_padded[i:i+8])
        plaintext += int_to_bytes(S[0] ^ Ci, 8)
        S[0] = Ci
        ascon_permutation(S, B)
    return plaintext + int_to_bytes(S[0], 8)[:c_lastlen]

def ascon_finalize(S, key):
    S[1] ^= bytes_to_int(key[0:8])
    S[2] ^= bytes_to_int(key[8:16])
    ascon_permutation(S, A)
    S[3] ^= bytes_to_int(key[0:8])
    S[4] ^= bytes_to_int(key[8:16])
    return int_to_bytes(S[3], 8) + int_to_bytes(S[4], 8)

def ascon_permutation(S, rounds):
    for r in range(12 - rounds, 12):
        S[2] ^= 0xf0 - r * 0x10 + r * 0x1
        S[0] ^= S[4]
        S[4] ^= S[3]
        S[2] ^= S[1]
        T = [(S[i] ^ 0xFFFFFFFFFFFFFFFF) & S[(i + 1) % 5] for i in range(5)]
        for i in range(5):
            S[i] ^= T[(i + 1) % 5]
        S[1] ^= S[0]
        S[0] ^= S[4]
        S[3] ^= S[2]
        S[2] ^= 0xFFFFFFFFFFFFFFFF
        S[0] ^= rotr(S[0], 19) ^ rotr(S[0], 28)
        S[1] ^= rotr(S[1], 61) ^ rotr(S[1], 39)
        S[2] ^= rotr(S[2], 1) ^ rotr(S[2], 6)
        S[3] ^= rotr(S[3], 10) ^ rotr(S[3], 17)
        S[4] ^= rotr(S[4], 7) ^ rotr(S[4], 41)

def bytes_to_int(bytes):
    return sum((b << (8 * i) for i, b in enumerate(reversed(bytes))))

def int_to_bytes(value, length):
    return bytes((value >> (8 * i) & 0xff for i in range(length - 1, -1, -1)))

def bytes_to_state(bytes):
    return [bytes_to_int(bytes[i * 8:(i + 1) * 8]) for i in range(5)]

def rotr(x, n):
    return (x >> n) | (x << (64 - n)) & 0xFFFFFFFFFFFFFFFF

def demo_aead():
    print("=== demo encryption using ascon128 ===")

    key   = b"\x00\x01\x02\x03\x04\x05\x06\x07\x08\x09\x0A\x0B\x0C\x0D\x0E\x0F"
    nonce = b"\x10\x11\x12\x13\x14\x15\x16\x17\x18\x19\x1A\x1B\x1C\x1D\x1E\x1F"
    associateddata = b"\x20\x21\x22\x23\x24\x25\x26\x27\x28\x29\x2A\x2B\x2C\x2D\x2E\x2F"
    plaintext      = b"\x30\x31\x32\x33\x34\x35\x36\x37\x38\x39\x3A\x3B\x3C\x3D\x3E\x3F"

    ciphertext        = ascon_aead_encrypt(key, nonce, associateddata, plaintext)
    receivedplaintext = ascon_aead_decrypt(key, nonce, associateddata, ciphertext)

    if receivedplaintext == None: print("verification failed!")
        
    demo_print([("key", key), 
                ("nonce", nonce), 
                ("plaintext", plaintext), 
                ("ass.data", associateddata), 
                ("ciphertext", ciphertext[:-16]), 
                ("tag", ciphertext[-16:]), 
                ("received", receivedplaintext), 
               ])

def demo_print(data):
    maxlen = max([len(text) for (text, val) in data])
    for text, val in data:
        print("{text}:{align} 0x{val} ({length} bytes)".format(text=text, align=((maxlen - len(text)) * " "), val=bytes_to_hex(val), length=len(val)))

def bytes_to_hex(b):
    return b.hex()

if __name__ == "__main__":
    demo_aead()